package com.example.tass.component;

import cn.hutool.core.io.resource.ClassPathResource;
import com.example.tass.exception.SignException;
import com.example.tass.licence.LicenceContent;
import com.example.tass.utils.IPUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.boot.ApplicationArguments;
import org.springframework.boot.ApplicationRunner;
import org.springframework.core.Ordered;
import org.springframework.stereotype.Component;

import java.io.*;
import java.net.SocketException;
import java.net.UnknownHostException;
import java.security.*;
import java.security.cert.Certificate;
import java.security.cert.CertificateException;
import java.security.cert.CertificateFactory;
import java.security.cert.X509Certificate;
import java.util.Base64;
import java.util.Date;
import java.util.List;

@Component
public class MyApplicationRunner implements ApplicationRunner, Ordered {

    private static final Logger log = LoggerFactory.getLogger(MyApplicationRunner.class);
    private static final String KEY_ALGORITHM_256 = "SHA256withRSA";

    // @Value(value = "${cerPath}")
    // private String cerPath;

    @Override
    public void run(ApplicationArguments args) throws Exception {
        String[] sourceArgs = args.getSourceArgs();

        String rootCert = null;
        String userCert = null;
        for (String sourceArg : sourceArgs) {
            if(sourceArg.contains("--rootCert=")){
                rootCert = sourceArg.substring("--rootCert=".length());
            }
            if(sourceArg.contains("--userCert=")){
                userCert = sourceArg.substring("--userCert=".length());
            }
        }

        // 读取序列化的licence文件，转换成LicenceContent对象
        ObjectInputStream objectInputStream = new ObjectInputStream(new ClassPathResource("licence.lic").getStream());
        LicenceContent licenceContent = (LicenceContent) objectInputStream.readObject();

        // 1. 校验用户证书
        checkCertificate(rootCert,userCert);

        // 2. 验证签名
        verifySign(licenceContent,userCert);

        // 3. 验证Licence是否在有效期内
        if (!validateDate(licenceContent)) {
            log.error("Expire: " + licenceContent.getNotAfterTime());
            throw new RuntimeException("License invalid");
        }

        // 4. 校验IP是否在白名单中
        if (!validateIp(licenceContent)) {
            log.error("IP error, only allow: " + licenceContent.getAllowIp());
            throw new RuntimeException("License invalid");
        }
    }

    private void verifySign(LicenceContent licenceContent,String userCertPath) throws GeneralSecurityException, IOException, SignException {
        // PublicKey publicKey = getPublicKey(userCertPath);
        PublicKey publicKey = getX509Certificate(userCertPath).getPublicKey();
        StringBuilder sb = new StringBuilder();
        sb.append(licenceContent.getUser())
                .append(licenceContent.getNotBeforeTime())
                .append(licenceContent.getNotAfterTime())
                .append(licenceContent.getAllowIp())
                .append(licenceContent.getMac());
        RSA rsa = new RSA();
        boolean verify = rsa.verify(sb.toString().getBytes(), publicKey, Base64.getDecoder().decode(licenceContent.getSign().getBytes()), KEY_ALGORITHM_256);
        log.info("The licence is valid...");
        if (!verify) {
            throw new RuntimeException("License invalid");
        }
    }

    private void checkCertificate(String rootCertPath,String userCertPath) throws CertificateException, IOException {
        if(rootCertPath == null || userCertPath == null){
            throw new RuntimeException("请提供根证书和用户证书...");
        }
        // 获取根证书
        X509Certificate x509RootCertificate = getX509Certificate(rootCertPath);
        // 获取用户证书
        X509Certificate x509UserCertificate = getX509Certificate(userCertPath);
        // 验证用户证书的有效性
        try {
            x509UserCertificate.verify(x509RootCertificate.getPublicKey());
            x509UserCertificate.checkValidity();
        } catch (NoSuchAlgorithmException | NoSuchProviderException | SignatureException | InvalidKeyException e) {
            throw new RuntimeException(e);
        }
        log.info("The user certificate is successfully verified...");
    }

    // 把Base64格式的pem文件转换成Certificate对象
    public X509Certificate getX509Certificate(String certificateName) throws IOException, CertificateException {
        // 获取根证书
        FileInputStream inputStream = new FileInputStream(certificateName);
        BufferedInputStream bufferedInputStream = new BufferedInputStream(inputStream);
        CertificateFactory certificateFactory = CertificateFactory.getInstance("X.509");
        X509Certificate certificate = null;
        while (bufferedInputStream.available() > 0) {
            certificate = (X509Certificate) certificateFactory.generateCertificate(bufferedInputStream);
            // System.out.println(certificate.toString());
        }
        return certificate;
    }

    public static PublicKey getPublicKey(String filename) throws IOException, GeneralSecurityException {
        String publicKeyPEM = getKey(filename);
        return getPublicKeyFromString(publicKeyPEM);
    }

    private static String getKey(String filename) throws IOException {
        // Read key from file
        String strKeyPEM = "";
        BufferedReader br = new BufferedReader(new FileReader(filename));
        String line;
        while ((line = br.readLine()) != null) {
            strKeyPEM += line;
        }
        br.close();
        return strKeyPEM;
    }

    public static PublicKey getPublicKeyFromString(String key) throws IOException, GeneralSecurityException {
        String publicKeyPEM = key;
        publicKeyPEM = publicKeyPEM.replace("-----BEGIN CERTIFICATE-----", "");
        publicKeyPEM = publicKeyPEM.replace("-----END CERTIFICATE-----", "");
        byte[] encoded = Base64.getDecoder().decode(publicKeyPEM);
        // KeyFactory kf = KeyFactory.getInstance("RSA");
        // RSAPublicKey pubKey = (RSAPublicKey) kf.generatePublic(new X509EncodedKeySpec(encoded));
        Certificate certificate = CertificateFactory.getInstance("X.509").generateCertificate(new ByteArrayInputStream(encoded));
        return certificate.getPublicKey();
    }

    // 校验时间有效期
    private static boolean validateDate(LicenceContent licenceContent) {
        Date currentTime = new Date();
        if (currentTime.after(licenceContent.getNotAfterTime())) {
            return false;
        } else {
            return true;
        }
    }

    // 校验ip白名单
    private static boolean validateIp(LicenceContent licenceContent) throws SocketException, UnknownHostException {
        List<String> localIps = IPUtils.getLocalIp();
        boolean hasAllowIp = localIps.contains(licenceContent.getAllowIp());
        if (!hasAllowIp) {
            log.warn("Local ip: " + (localIps.size() > 0 ? localIps.get(0) : "fail to find ip"));
        }
        return hasAllowIp;
    }

    @Override
    public int getOrder() {
        return Ordered.HIGHEST_PRECEDENCE;
    }
}
